import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import csv
# df = pd.read_csv('./files/平安银行2.csv', index_col='date', parse_dates=['date'])[['open','close','low','high','date']]
df = pd.read_csv('./files/通威股份.csv')[['open','close','low','high','date']]
df['date'] = pd.to_datetime(df['date'])
df['ma5'] = np.nan
df['ma30'] = np.nan

# 设定你想要筛选的日期范围
# start_date = '2021-01-02'
# end_date = '2017-12-15'
end_date = '2024-7-18'

 
# 使用布尔索引进行筛选
# mask = (df['date'] >= start_date) & (df['date'] <= end_date)
mask = df['date'] <= end_date
df = df.loc[mask]
# print(df)

for i in range(4, len(df)):
  df.loc[df.index[i], 'ma5'] = df['close'][i-4:i+1].mean()

for i in range(29, len(df)):
  df.loc[df.index[i], 'ma30'] = df['close'][i-29:i+1].mean()

# 简单方法求均线
df['ma5_T'] = df['close'].rolling(window=5).mean()
df['ma30_T'] = df['close'].rolling(window=30).mean()

print(df)
# df[['close','ma5','ma30']].plot()
# plt.show()

golden_cross = []
death_cross = []

for i in range(1,len(df)):
  if df['ma5'][i] >= df['ma30'][i] and df['ma5'][i-1] < df['ma30'][i-1]:
    golden_cross.append(df['date'][i].to_pydatetime())
  if df['ma5'][i] <= df['ma30'][i] and df['ma5'][i-1] > df['ma30'][i-1]:
    death_cross.append(df['date'][i].to_pydatetime())


# 方式2   需要索引改日期
# sr1 = df['ma5'] < df['ma30']
# sr2 = df['ma5'] >= df['ma30']

# death_cross = df[sr1 & sr2.shift(1)].shift
# golden_cross = df[~(sr1 | sr2.shift(1))].index


print(death_cross)
print(golden_cross)


filename = './files/通威股份金叉和死叉.csv'
mode ='w'
with open(filename, mode, newline='') as file:
  writer = csv.writer(file)
  for de, go in zip(death_cross, golden_cross):  
    writer.writerow([de, go])

file.close()
